package com.yifeng.repo.micro.service.server.engine.routing;

import java.util.*;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Java实现通过权重实现随机
 * <p>
 * 【需求背景】从不同权重的N个元素中随机选择一个，并使得总体选择结果是按照权重分布的。
 * 【实现思路】比如我们现在有三台服务器(IP1、IP2、IP3)，权重分别为1，3，2。现在想对三台服务器 做负载均衡，可见：
 *           权重总和 =  三台机器权重相加
 *           权重比例 = 自己的权重 / 总权重
 *              IP1           IP2           IP3
 *            weight      weight           weight
 *               1             3             2
 *              radio       radio         radio
 *              1/6          3/6           2/6
 *           使用随机数函数(ThreadLocalRandom 或者 Random )，取 [0,6] 之间的随机数，根据随机数落在哪个范围决定如何选择。
 *           例如随机数为 2，处于 [1,4] 范围，那么就选择 IP2。
 * <p>
 * https://blog.csdn.net/qq_40083897/article/details/126969186
 */
public class WeightRandom<T> {
    private final List<T> items = new ArrayList<>();
    private double[] weights;
    private final long timestamp;

    public WeightRandom(List<ItemWithWeight<T>> itemsWithWeight) {
        this.calWeights(itemsWithWeight);
        this.timestamp = System.currentTimeMillis();
    }

    /**
     * 计算权重，初始化或者重新定义权重时使用
     */
    public void calWeights(List<ItemWithWeight<T>> itemsWithWeight) {
        items.clear();

        // 计算权重总和
        double originWeightSum = 0;
        for (ItemWithWeight<T> itemWithWeight : itemsWithWeight) {
            double weight = itemWithWeight.getWeight();
            if (weight <= 0) {
                continue;
            }

            items.add(itemWithWeight.getItem());
            if (Double.isInfinite(weight)) {
                weight = 10000.0D;
            }
            if (Double.isNaN(weight)) {
                weight = 1.0D;
            }
            originWeightSum += weight;
        }

        // 计算每个item的实际权重比例
        double[] actualWeightRatios = new double[items.size()];
        int index = 0;
        for (ItemWithWeight<T> itemWithWeight : itemsWithWeight) {
            double weight = itemWithWeight.getWeight();
            if (weight <= 0) {
                continue;
            }
            actualWeightRatios[index++] = weight / originWeightSum;
        }

        // 计算每个item的权重范围
        // 权重范围起始位置
        weights = new double[items.size()];
        double weightRangeStartPos = 0;
        for (int i = 0; i < index; i++) {
            weights[i] = weightRangeStartPos + actualWeightRatios[i];
            weightRangeStartPos += actualWeightRatios[i];
        }
    }

    /**
     * 基于权重随机算法选择
     */
    public T choose() {
        double random = ThreadLocalRandom.current().nextDouble();
        int index = Arrays.binarySearch(weights, random);
        if (index < 0) {
            index = -index - 1;
        } else {
            return items.get(index);
        }

        if (index < weights.length && random < weights[index]) {
            return items.get(index);
        }

        // 通常不会走到这里，为了保证能得到正确的返回，这里随便返回一个
        return items.get(0);
    }

    public static class ItemWithWeight<T> {
        T item;
        double weight;

        public ItemWithWeight() {
        }

        public ItemWithWeight(T item, double weight) {
            this.item = item;
            this.weight = weight;
        }

        public T getItem() {
            return item;
        }

        public void setItem(T item) {
            this.item = item;
        }

        public double getWeight() {
            return weight;
        }

        public void setWeight(double weight) {
            this.weight = weight;
        }
    }

    public long getTimestamp() {
        return timestamp;
    }

    public static void main(String[] args) {
        int sampleCount = 100;

        ItemWithWeight<String> server1 = new ItemWithWeight<>("ip1", 1.0);
        ItemWithWeight<String> server2 = new ItemWithWeight<>("ip2", 2.0);
        ItemWithWeight<String> server3 = new ItemWithWeight<>("ip3", 3.0);

        WeightRandom<String> weightRandom = new WeightRandom<>(Arrays.asList(server1, server2, server3));

        Map<String, AtomicInteger> statistics = new HashMap<>();
        for (int i = 0; i < sampleCount; i++) {
            statistics.computeIfAbsent(weightRandom.choose(), (k) -> new AtomicInteger()).incrementAndGet();
        }
        statistics.forEach((k, v) -> {
            double hit = (double) v.get() / sampleCount;
            System.out.println(k + ", count:" + v.get() + ", hit:" + hit);
        });
    }
}